#

#from InsightToolkit import *
import itk

from sys import argv
from math import pi

class reg():

    def __init__(self):
        """
        """
        self._fixedImageReader  = itk.ImageFileReader.IF2.New()
        self._movingImageReader = itk.ImageFileReader.IF2.New()
        self._imageMetric     = itk.MeanSquaresImageToImageMetric.IF2IF2.New()

        self._firstRun = True

    def SetFixedFileName(self,fixedImageFileName):
        """
        should only be required once (since the image is fixed, right)
        """
        self._fixedImageReader.SetFileName(  fixedImageFileName )
        self._fixedImageReader.Update()
        self._fixedImage  = self._fixedImageReader.GetOutput()


    def SetMovingFileName(self,movingImageFileName):
        """
        """
        self._movingImageFileName = movingImageFileName
        self._movingImageReader.SetFileName( self._movingImageFileName )
        self._movingImage = self._movingImageReader.GetOutput()


    def UpdateMovingImage(self):
        """
        Warning: This will only work if the moving image is consistently saved
        to the same filename!
        """
        self._movingImageReader.Update()
        self._movingImage = self._movingImageReader.GetOutput()

    def SetImageMetric(self,metric):
        """
        Allow the user to specify the desired ITK metric, current choices:
            MS = MeanSquares [Default]
            MMI = MattesMutualInformation
            NC = NormalizedCorrelation
            MRSD = MeanReciprocalSquareDifference
            MI = MutualInformation
            GD = GradientDifference
        """
        if metric.lower() == 'ms':
            self._imageMetric     = itk.MeanSquaresImageToImageMetric.IF2IF2.New()
        if metric.lower() == 'mmi':
            self._imageMetric     = itk.MattesMutualInformationImageToImageMetric.IF2IF2.New()
        if metric.lower() == 'nc':
            self._imageMetric     = itk.NormalizedCorrelationImageToImageMetric.IF2IF2.New()
        if metric.lower() == 'mrsd':
            self._imageMetric     = itk.MeanReciprocalSquareDifferenceImageToImageMetric.IF2IF2.New()
        if metric.lower() == 'mi':
            self._imageMetric     = itk.MutualInformationImageToImageMetric.IF2IF2.New()
        if metric.lower() == 'gd':
            self._imageMetric     = itk.GradientDifferenceImageToImageMetric.IF2IF2.New()


    def GetMetric(self):
        #
        #  Instantiate the classes for the registration framework
        #
        self._registration    = itk.ImageRegistrationMethod.IF2IF2.New()
        # Metric has already been set
        self._transform       = itk.CenteredRigid2DTransform.D.New()
        self._optimizer       = itk.RegularStepGradientDescentOptimizer.New()
        self._interpolator    = itk.LinearInterpolateImageFunction.IF2D.New()

        self._registration.SetOptimizer(      self._optimizer.GetPointer() )
        self._registration.SetTransform(      self._transform.GetPointer() )
        self._registration.SetInterpolator(   self._interpolator.GetPointer() )
        self._registration.SetMetric(         self._imageMetric.GetPointer() )

        self._registration.SetFixedImage(  self._fixedImage )
        self._registration.SetFixedImageRegion(  self._fixedImage.GetBufferedRegion() )
        self._imageMetric.SetFixedImage( self._fixedImage )
        self._imageMetric.SetMovingImage( self._movingImage )

        self._initialParameters = self._transform.GetParameters()
        self._initialParameters

        self._optimizer.SetNumberOfIterations( 1 )

        self.UpdateMovingImage()
        self._registration.SetMovingImage( self._movingImage )


        self._initialParameters.fill(0)

        self._registration.SetInitialTransformParameters( self._initialParameters )
        self._transform.SetParameters( self._initialParameters )

        self._registration.StartRegistration()

        return self._imageMetric.GetValue(self._initialParameters)


if __name__ == "__main__":
    fixedImageFileName = argv[1]
    movingImageFileName = argv[2]
    reg=reg()
    reg.GetMetric(fixedImageFileName, movingImageFileName)